import json
import requests
import threading
import time
import concurrent.futures
from collections import defaultdict

# 服务器配置
SERVER_URL = "http://8.152.3.5:5000/auto_test"  # 根据实际情况修改

# 测试参数
MAX_WORKERS = 3  # 固定5个并发
TEST_DURATION = 2  # 测试持续时间(秒)
REQUEST_TIMEOUT = 300  # 请求超时时间(秒)


def send_request(request_id):
    """
    发送POST请求到/auto_test端点，正确地通过multipart/form-data上传文件
    """
    url = "http://8.152.3.5:5000/auto_test"  # 根据实际情况修改URL

    # 准备JSON数据（作为表单字段发送）
    data = {
        'brand': 'mi',
        'testname': 'test1',
        'case': 'test_case5',
        'rpk_name': '知识大百科',
        'entity': 'hn'
    }

    # 准备文件
    file_path = r"C:\Users\chai0\sns-project\Analyze-advertising-events\Determining_ad\test\test\historical_data\data_hn_20250826_161828.parquet"

    start_time = time.time()
    try:
        with open(file_path, 'rb') as f:
            # 正确地通过multipart/form-data发送文件和数据
            files = {'all_df': f}
            response = requests.post(url, data=data, files=files, timeout=REQUEST_TIMEOUT)

        response_time = time.time() - start_time

        result = {
            'request_id': request_id,
            'success': True,
            'status_code': response.status_code,
            'response_time': response_time,
            'response_text': response.text
        }

        print(f"请求 {request_id} 完成: 状态码 {response.status_code}, 响应时间 {response_time:.2f} 秒")

    except FileNotFoundError:
        response_time = time.time() - start_time
        result = {
            'request_id': request_id,
            'success': False,
            'error': f"文件未找到: {file_path}",
            'response_time': response_time
        }
        print(f"请求 {request_id} 失败: 文件未找到")
    except requests.exceptions.Timeout:
        response_time = time.time() - start_time
        result = {
            'request_id': request_id,
            'success': False,
            'error': 'Request timeout',
            'response_time': REQUEST_TIMEOUT
        }
        print(f"请求 {request_id} 超时")
    except Exception as e:
        response_time = time.time() - start_time
        result = {
            'request_id': request_id,
            'success': False,
            'error': str(e),
            'response_time': response_time
        }
        print(f"请求 {request_id} 失败: {e}")

    return result


def run_concurrent_test(num_workers, duration):
    """运行并发测试"""
    print(f"开始并发测试: {num_workers} 个并发请求，持续 {duration} 秒")

    start_time = time.time()
    end_time = start_time + duration
    results = []
    request_count = 0
    last_print_time = start_time

    with concurrent.futures.ThreadPoolExecutor(max_workers=num_workers) as executor:
        futures = []

        # 在指定时间内持续发送请求
        while time.time() < end_time:
            try:
                future = executor.submit(send_request, request_count)
                futures.append(future)
                request_count += 1
            except concurrent.futures.TimeoutError:
                print("提交请求超时，跳过")

            # 控制请求发送速率
            time.sleep(0.3)

            # 每5秒打印一次进度
            current_time = time.time()
            if current_time - last_print_time >= 5:
                elapsed = current_time - start_time
                completed = sum(1 for f in futures if f.done())
                print(f"[进度] {elapsed:.1f}/{duration}s | 已提交: {request_count} | 已完成: {completed}")
                last_print_time = current_time

        # 等待请求完成，但不超过额外的10秒
        completed_results = []
        hard_deadline = time.time() + 10

        # 只收集在截止时间前完成的请求
        for future in futures:
            if time.time() >= hard_deadline:
                print("达到硬性截止时间，停止收集结果")
                break

            if future.done():
                try:
                    result = future.result()
                    completed_results.append(result)
                except Exception as e:
                    print(f"请求处理异常: {e}")
                    completed_results.append({
                        'request_id': 'unknown',
                        'success': False,
                        'error': str(e),
                        'response_time': REQUEST_TIMEOUT
                    })
            else:
                # 取消未完成的请求
                future.cancel()


        # 统计未完成请求数
        unfinished_count = len(futures) - len(completed_results)
        if unfinished_count > 0:
            print(f"有 {unfinished_count} 个请求未完成，已取消")

        results = completed_results

    # 分析结果
    successful_requests = [r for r in results if r.get('success', False)]
    failed_requests = [r for r in results if not r.get('success', False)]

    total_requests = len(results)
    successful_count = len(successful_requests)
    failed_count = len(failed_requests)

    avg_response_time = sum([r['response_time'] for r in successful_requests]) / len(successful_requests) if successful_requests else 0
    max_response_time = max([r['response_time'] for r in successful_requests]) if successful_requests else 0
    min_response_time = min([r['response_time'] for r in successful_requests]) if successful_requests else 0

    # 计算每秒请求数
    test_duration_actual = min(time.time() - start_time, duration + 10)  # 最大不超过设定时间+10秒
    rps = total_requests / test_duration_actual if test_duration_actual > 0 else 0

    print(f"   测试完成: {num_workers} 并发")
    print(f"  总请求数: {total_requests}")
    print(f"  成功请求数: {successful_count}")
    print(f"  失败请求数: {failed_count}")
    print(f"  成功率: {successful_count / total_requests * 100:.2f}%") if total_requests > 0 else print("  成功率: 0%")
    print(f"  平均响应时间: {avg_response_time:.2f} 秒")
    print(f"  最大响应时间: {max_response_time:.2f} 秒")
    print(f"  最小响应时间: {min_response_time:.2f} 秒")
    print(f"  每秒请求数: {rps:.2f}")
    print(f"  实际测试时间: {test_duration_actual:.2f} 秒")
    print("-" * 50)

    return {
        'total_requests': total_requests,
        'successful_requests': successful_count,
        'failed_requests': failed_count,
        'success_rate': successful_count / total_requests if total_requests > 0 else 0,
        'avg_response_time': avg_response_time,
        'max_response_time': max_response_time,
        'min_response_time': min_response_time,
        'rps': rps,
        'duration': test_duration_actual,
        'results': results
    }


def main():
    """主函数"""
    print("开始压力测试...")
    print(f"目标服务器: {SERVER_URL}")
    print(f"并发数: {MAX_WORKERS}")
    print(f"测试持续时间: {TEST_DURATION} 秒")
    print("=" * 60)

    # 运行测试
    try:
        result = run_concurrent_test(MAX_WORKERS, TEST_DURATION)
        print("\n压力测试完成!")
        return result
    except KeyboardInterrupt:
        print("\n测试被用户中断")
        return None
    except Exception as e:
        print(f"测试过程中出现错误: {e}")
        return None


# 运行这个程序
if __name__ == "__main__":
    main()
